import torch
import w2v_costh


x_norm_fwd = torch.randn(3, 768).cuda(0)
w_fwd = torch.randn(768, 1211).cuda(0)

x_norm = torch.randn(3, 768).cuda(0)
w = torch.randn(768, 1211).cuda(0)
grad_mm = torch.randn(3, 1211).cuda(0)
costh_out = torch.randn(768, 1211).cuda(0)
norm_vals = torch.randn(1211, ).cuda(0)

bkwd_res = w2v_costh.backward(grad_mm, x_norm, w, costh_out, norm_vals)
fwd_res = w2v_costh.forward(x_norm_fwd, w_fwd)

print(f"Length of backward return value is {len(bkwd_res)}")
print("The shape of backward return value[0] is: ", bkwd_res[0].shape, bkwd_res[0].device)
print("The shape of backward return value[1] is: ", bkwd_res[1].shape, bkwd_res[1].device)

print(f"Length of forward return value is {len(fwd_res)}")
print("The shape of forward return value[0] is: ", fwd_res[0].shape, fwd_res[0].device)
print("The shape of forward return value[1] is: ", fwd_res[1].shape, fwd_res[1].device)
print("The shape of forward return value[2] is: ", fwd_res[2].shape, fwd_res[2].device)

